#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
知识库管理器
整合知识库构建、搜索、管理功能
"""

import os
import json
import argparse
import logging
from typing import Dict, Any, List
from knowledge_base_builder import KnowledgeBaseBuilder
from knowledge_base_search import KnowledgeBaseSearch

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class KnowledgeBaseManager:
    """知识库管理器"""
    
    def __init__(self, kb_path: str = "knowledge_base.json"):
        """
        初始化知识库管理器
        
        Args:
            kb_path: 知识库文件路径
        """
        self.kb_path = kb_path
        self.builder = None
        self.search = None
        
    def build_knowledge_base(self, pdf_path: str = "论人与自然的相处之道.pdf") -> Dict[str, Any]:
        """
        构建知识库
        
        Args:
            pdf_path: PDF文件路径
            
        Returns:
            Dict[str, Any]: 构建结果
        """
        logger.info("开始构建知识库...")
        
        try:
            # 创建构建器
            self.builder = KnowledgeBaseBuilder(pdf_path)
            
            # 构建知识库
            knowledge_base = self.builder.build_knowledge_base()
            
            # 保存知识库
            output_path = self.builder.save_knowledge_base(self.kb_path)
            
            # 初始化搜索引擎
            self.search = KnowledgeBaseSearch(self.kb_path)
            
            result = {
                "status": "success",
                "message": "知识库构建完成",
                "output_path": output_path,
                "statistics": {
                    "total_entities": knowledge_base["metadata"]["total_entities"],
                    "total_relations": knowledge_base["metadata"]["total_relations"],
                    "total_chunks": knowledge_base["metadata"]["total_chunks"],
                    "total_characters": knowledge_base["metadata"]["total_characters"]
                }
            }
            
            logger.info("知识库构建完成")
            return result
            
        except Exception as e:
            logger.error(f"构建知识库失败: {e}")
            return {
                "status": "error",
                "message": f"构建失败: {str(e)}"
            }
    
    def load_knowledge_base(self) -> bool:
        """
        加载知识库
        
        Returns:
            bool: 加载是否成功
        """
        try:
            if not os.path.exists(self.kb_path):
                logger.warning(f"知识库文件不存在: {self.kb_path}")
                return False
            
            self.search = KnowledgeBaseSearch(self.kb_path)
            logger.info("知识库加载成功")
            return True
            
        except Exception as e:
            logger.error(f"加载知识库失败: {e}")
            return False
    
    def search_knowledge_base(self, query: str, search_type: str = "all", top_k: int = 10) -> Dict[str, Any]:
        """
        搜索知识库
        
        Args:
            query: 搜索查询
            search_type: 搜索类型 (all, entities, relations, chunks)
            top_k: 返回结果数量
            
        Returns:
            Dict[str, Any]: 搜索结果
        """
        if not self.search:
            if not self.load_knowledge_base():
                return {
                    "status": "error",
                    "message": "知识库未加载"
                }
        
        try:
            if search_type == "entities":
                results = self.search.search_entities(query)[:top_k]
            elif search_type == "relations":
                results = self.search.search_relations(query)[:top_k]
            elif search_type == "chunks":
                results = self.search.search_chunks(query, top_k)
            else:  # all
                results = self.search.search_all(query, top_k)
            
            # 格式化结果
            if search_type == "all":
                formatted_results = {
                    "entities": [
                        {
                            "id": result.id,
                            "title": result.title,
                            "content": result.content,
                            "score": result.score,
                            "metadata": result.metadata
                        }
                        for result in results["entities"]
                    ],
                    "relations": [
                        {
                            "id": result.id,
                            "title": result.title,
                            "content": result.content,
                            "score": result.score,
                            "metadata": result.metadata
                        }
                        for result in results["relations"]
                    ],
                    "chunks": [
                        {
                            "id": result.id,
                            "title": result.title,
                            "content": result.content,
                            "score": result.score,
                            "metadata": result.metadata
                        }
                        for result in results["chunks"]
                    ]
                }
            else:
                formatted_results = [
                    {
                        "id": result.id,
                        "title": result.title,
                        "content": result.content,
                        "score": result.score,
                        "metadata": result.metadata
                    }
                    for result in results
                ]
            
            return {
                "status": "success",
                "query": query,
                "search_type": search_type,
                "results": formatted_results,
                "total_results": len(formatted_results) if search_type != 'all' else sum(len(v) for v in formatted_results.values())
            }
            
        except Exception as e:
            logger.error(f"搜索失败: {e}")
            return {
                "status": "error",
                "message": f"搜索失败: {str(e)}"
            }
    
    def get_entity_info(self, entity_name: str) -> Dict[str, Any]:
        """
        获取实体信息
        
        Args:
            entity_name: 实体名称
            
        Returns:
            Dict[str, Any]: 实体信息
        """
        if not self.search:
            if not self.load_knowledge_base():
                return {
                    "status": "error",
                    "message": "知识库未加载"
                }
        
        try:
            # 搜索实体
            entities = self.search.search_entities(entity_name)
            
            if not entities:
                return {
                    "status": "error",
                    "message": f"未找到实体: {entity_name}"
                }
            
            entity = entities[0]  # 取最匹配的实体
            
            # 获取相关关系
            relations = self.search.get_entity_relations(entity_name)
            
            # 获取相关实体
            related_entities = self.search.get_related_entities(entity_name)
            
            return {
                "status": "success",
                "entity": {
                    "id": entity.id,
                    "name": entity.title,
                    "type": entity.metadata["entity_type"],
                    "description": entity.content,
                    "keywords": entity.metadata["keywords"],
                    "confidence": entity.metadata["confidence"]
                },
                "relations": relations,
                "related_entities": related_entities
            }
            
        except Exception as e:
            logger.error(f"获取实体信息失败: {e}")
            return {
                "status": "error",
                "message": f"获取实体信息失败: {str(e)}"
            }
    
    def get_knowledge_graph(self, entity_name: str, depth: int = 2) -> Dict[str, Any]:
        """
        获取知识图谱
        
        Args:
            entity_name: 中心实体名称
            depth: 图谱深度
            
        Returns:
            Dict[str, Any]: 知识图谱
        """
        if not self.search:
            if not self.load_knowledge_base():
                return {
                    "status": "error",
                    "message": "知识库未加载"
                }
        
        try:
            graph = self.search.get_knowledge_graph(entity_name, depth)
            
            return {
                "status": "success",
                "center_entity": entity_name,
                "depth": depth,
                "graph": graph
            }
            
        except Exception as e:
            logger.error(f"获取知识图谱失败: {e}")
            return {
                "status": "error",
                "message": f"获取知识图谱失败: {str(e)}"
            }
    
    def get_statistics(self) -> Dict[str, Any]:
        """
        获取知识库统计信息
        
        Returns:
            Dict[str, Any]: 统计信息
        """
        if not self.search:
            if not self.load_knowledge_base():
                return {
                    "status": "error",
                    "message": "知识库未加载"
                }
        
        try:
            stats = self.search.get_statistics()
            
            return {
                "status": "success",
                "statistics": stats
            }
            
        except Exception as e:
            logger.error(f"获取统计信息失败: {e}")
            return {
                "status": "error",
                "message": f"获取统计信息失败: {str(e)}"
            }
    
    def interactive_search(self):
        """交互式搜索"""
        if not self.load_knowledge_base():
            print("❌ 知识库加载失败")
            return
        
        print("=" * 60)
        print("🔍 知识库交互式搜索")
        print("=" * 60)
        print("输入 'quit' 或 'exit' 退出")
        print("输入 'help' 查看帮助")
        print("输入 'stats' 查看统计信息")
        print("-" * 60)
        
        while True:
            try:
                query = input("\n💬 请输入搜索查询: ").strip()
                
                if query.lower() in ['quit', 'exit', '退出']:
                    print("👋 再见！")
                    break
                
                if query.lower() in ['help', '帮助']:
                    print("\n📖 帮助信息:")
                    print("  - 可以搜索实体、关系、知识块")
                    print("  - 例如: '人与自然'、'生态危机'、'可持续发展'")
                    print("  - 输入 'stats' 查看知识库统计信息")
                    print("  - 输入 'quit' 或 'exit' 退出")
                    continue
                
                if query.lower() in ['stats', '统计']:
                    stats_result = self.get_statistics()
                    if stats_result['status'] == 'success':
                        stats = stats_result['statistics']
                        print(f"\n📊 知识库统计信息:")
                        print(f"  总实体数: {stats['total_entities']}")
                        print(f"  总关系数: {stats['total_relations']}")
                        print(f"  总知识块数: {stats['total_chunks']}")
                        
                        print(f"\n实体类型分布:")
                        for entity_type, count in stats['entity_types'].items():
                            print(f"    {entity_type}: {count}")
                        
                        print(f"\n关系类型分布:")
                        for relation_type, count in stats['relation_types'].items():
                            print(f"    {relation_type}: {count}")
                    continue
                
                if not query:
                    print("❌ 查询不能为空，请重新输入")
                    continue
                
                print(f"\n🔍 正在搜索: {query}")
                print("-" * 40)
                
                # 执行搜索
                result = self.search_knowledge_base(query, search_type="all", top_k=5)
                
                if result['status'] == 'success':
                    print(f"✅ 找到 {result['total_results']} 个结果")
                    
                    # 显示实体结果
                    if result['results']['entities']:
                        print(f"\n📋 实体结果:")
                        for i, entity in enumerate(result['results']['entities'], 1):
                            print(f"  {i}. {entity['title']} (分数: {entity['score']:.2f})")
                            print(f"     类型: {entity['metadata']['entity_type']}")
                            print(f"     描述: {entity['content'][:100]}...")
                    
                    # 显示关系结果
                    if result['results']['relations']:
                        print(f"\n🔗 关系结果:")
                        for i, relation in enumerate(result['results']['relations'], 1):
                            print(f"  {i}. {relation['title']} (分数: {relation['score']:.2f})")
                            print(f"     描述: {relation['content'][:100]}...")
                    
                    # 显示知识块结果
                    if result['results']['chunks']:
                        print(f"\n📄 知识块结果:")
                        for i, chunk in enumerate(result['results']['chunks'], 1):
                            print(f"  {i}. {chunk['title']} (分数: {chunk['score']:.2f})")
                            print(f"     内容: {chunk['content'][:100]}...")
                else:
                    print(f"❌ 搜索失败: {result['message']}")
                
            except KeyboardInterrupt:
                print("\n👋 再见！")
                break
            except Exception as e:
                print(f"❌ 发生错误: {e}")

def main():
    """主函数"""
    parser = argparse.ArgumentParser(description="知识库管理器")
    parser.add_argument("--mode", choices=["build", "search", "interactive", "stats"], 
                       default="interactive", help="运行模式")
    parser.add_argument("--pdf", default="论人与自然的相处之道.pdf", help="PDF文件路径")
    parser.add_argument("--query", help="搜索查询")
    parser.add_argument("--entity", help="实体名称")
    parser.add_argument("--graph", help="知识图谱中心实体")
    parser.add_argument("--depth", type=int, default=2, help="图谱深度")
    parser.add_argument("--kb-path", default="knowledge_base.json", help="知识库文件路径")
    
    args = parser.parse_args()
    
    # 创建知识库管理器
    manager = KnowledgeBaseManager(args.kb_path)
    
    try:
        if args.mode == "build":
            # 构建知识库
            result = manager.build_knowledge_base(args.pdf)
            print(f"构建结果: {result}")
            
        elif args.mode == "search":
            # 搜索知识库
            if not args.query:
                print("❌ 搜索模式需要指定 --query 参数")
                return
            
            result = manager.search_knowledge_base(args.query)
            print(f"搜索结果: {result}")
            
        elif args.mode == "interactive":
            # 交互式搜索
            manager.interactive_search()
            
        elif args.mode == "stats":
            # 显示统计信息
            result = manager.get_statistics()
            if result['status'] == 'success':
                stats = result['statistics']
                print(f"知识库统计信息:")
                print(f"  总实体数: {stats['total_entities']}")
                print(f"  总关系数: {stats['total_relations']}")
                print(f"  总知识块数: {stats['total_chunks']}")
                
                print(f"\n实体类型分布:")
                for entity_type, count in stats['entity_types'].items():
                    print(f"    {entity_type}: {count}")
                
                print(f"\n关系类型分布:")
                for relation_type, count in stats['relation_types'].items():
                    print(f"    {relation_type}: {count}")
            else:
                print(f"获取统计信息失败: {result['message']}")
        
        # 额外的实体信息查询
        if args.entity:
            result = manager.get_entity_info(args.entity)
            print(f"实体信息: {result}")
        
        # 额外的知识图谱查询
        if args.graph:
            result = manager.get_knowledge_graph(args.graph, args.depth)
            print(f"知识图谱: {result}")
            
    except KeyboardInterrupt:
        print("\n👋 程序已停止")
    except Exception as e:
        logger.error(f"程序运行失败: {e}")
        raise

if __name__ == "__main__":
    main()
